import itertools

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv
import numpy as np
# from functorch import vmap
from torch import vmap


class GlobalGNN(nn.Module):
    def __init__(self, in_features, hidden_features):
        super().__init__()
        self.conv1 = GATConv(in_features, hidden_features)
        self.conv2 = GATConv(hidden_features, hidden_features)
        
    def forward(self, x):
        device = next(self.parameters()).device
        groups = x.shape[0]
        indices = list(range(groups))
        edge_index = torch.tensor(list(itertools.combinations(range(len(indices)), 2)), 
                                  dtype=torch.long).t().contiguous().to(device)
        out, (idx, attr) = self.conv1(x.to(torch.float32), edge_index, return_attention_weights=True)
        out = F.relu(out)
        out = self.conv2(out, idx, attr)
        
        return out


class LocalGNN(nn.Module):
    def __init__(self, num_node_features, hidden_features):
        super().__init__()
        self.conv1 = GATConv(num_node_features, hidden_features)
        self.conv2 = GATConv(hidden_features, hidden_features)
    
    def forward(self, x, edge_attr_raw):

        device = next(self.parameters()).device
        sites = edge_attr_raw.shape[0]
        indices = list(range(sites))
        edge_index = torch.tensor(list(itertools.combinations(range(len(indices)), 2)), 
                                  dtype=torch.long).t().contiguous().to(device)
        edge_attr = self._get_edge_attr(edge_attr_raw).to(device)
        out, (idx, attr) = self.conv1(x, edge_index, edge_attr, return_attention_weights=True)
        out = F.relu(out)
        out = self.conv2(out, idx, attr)
        
        return out

    def _get_edge_attr(self, edge_attr_raw):
        sites = edge_attr_raw.shape[0]
        edge_distances = []
        for i, j in itertools.combinations(range(sites), 2):
            lon1, lat1 = edge_attr_raw[i]
            lon2, lat2 = edge_attr_raw[j]

            lon1, lat1, lon2, lat2 = map(torch.deg2rad, [lon1, lat1, lon2, lat2])

            dlon = lon2 - lon1
            dlat = lat2 - lat1
            a = torch.sin(dlat / 2.0)**2 + torch.cos(lat1) * torch.cos(lat2) * torch.sin(dlon / 2.0)**2
            c = 2 * torch.arcsin(torch.sqrt(a))
            r = 6371  # r of earth (km)
            distance = c * r
            
            edge_distances.append(distance)
        edge_attr = torch.tensor([1.0 / d for d in edge_distances], dtype=torch.float)
        return edge_attr


class BatchCombinedModel(nn.Module):
    def __init__(self, local_in, local_out, global_in, global_out):
        super().__init__()
        self.local_model = LocalGNN(local_in, local_out)
        self.global_model = GlobalGNN(global_in, global_out)

    def forward(self, tot_x, coord_info):
        local_results = {}
        device = next(self.local_gnn.parameters()).device

        for key in tot_x.keys():
            x = tot_x[key].to(device)
            coord = torch.tensor(coord_info[key], dtype=torch.float, device=device) 
            B, T, N_k, H = x.shape

            x_flat = x.reshape(-1, N_k, H)

            def run_gnn_single(x_i):
                return self.local_gnn(x_i, coord)
            local_pred = vmap(run_gnn_single)(x_flat)

            local_pred = local_pred.reshape(B, T, N_k, -1)
            local_results[key] = local_pred

            local_total = torch.cat([local_results[k] for k in local_results.keys()], dim=-2)

            local_means = [torch.mean(local_results[k], dim=-2, keepdim=True) for k in local_results.keys()]
            local_mean = torch.cat(local_means, dim=-2)


class BatchCombinedModel(nn.Module):
    def __init__(self, local_in=1, local_out=64, global_in=64, global_out=64):
        super().__init__()
        self.local_gnn = LocalGNN(local_in, local_out)
        self.global_gnn = GlobalGNN(global_in, global_out)

    def forward(self, tot_x, coord_info):

        local_results = {}
        device = next(self.local_gnn.parameters()).device

        for key in tot_x.keys():
            x = tot_x[key].to(device)
            coord = torch.tensor(coord_info[key], dtype=torch.float, device=device)
            B, T, N_k, H = x.shape

            x_flat = x.reshape(-1, N_k, H)

            def run_gnn_single(x_i):
                return self.local_gnn(x_i, coord) 


            local_pred = vmap(run_gnn_single)(x_flat)

            local_pred = local_pred.reshape(B, T, N_k, -1) 
            local_results[key] = local_pred

        local_total = torch.cat([local_results[k] for k in local_results.keys()], dim=-2)

        local_means = [torch.mean(local_results[k], dim=-2, keepdim=True) for k in local_results.keys()]
        local_mean = torch.cat(local_means, dim=-2) 


        B, T, R, D = local_mean.shape
        local_mean_flat = local_mean.reshape(-1, R, D) 

        def run_global_gnn(x_g):
            return self.global_gnn(x_g) 

        global_output_flat = vmap(run_global_gnn)(local_mean_flat)
        global_output = global_output_flat.reshape(B, T, R, -1)
        
        return global_output, local_total
